import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from pathlib import Path
import random

# 数据集根目录
root_dir = "D:/document/post/datasets/GF7/Images and Shpfiles"
# 城市列表
cities = ["Chongqing", "Guangzhou", "Lanzhou", "Ningbo", "Shenzhen", "Tianjin"]
# 划分比例
train_ratio = 0.7
val_ratio = 0.15
test_ratio = 0.15

# 保存路径
output_dir = "D:/document/post/datasets/GF7/Processed"

# 处理每个城市
for city in cities:
    print(f"处理城市：{city}")
    city_path = os.path.join(root_dir, city)

    # 创建保存目录
    save_dir = os.path.join(output_dir, city)
    os.makedirs(save_dir, exist_ok=True)

    # 获取该城市下的所有区域
    areas = [area for area in os.listdir(city_path) if os.path.isdir(os.path.join(city_path, area))]

    # 所有数据文件路径
    image_paths = []
    label_paths = []

    # 遍历每个区域
    for area in areas:
        print(f"处理区域：{area}")
        area_path = os.path.join(city_path, area)
        image_path = os.path.join(area_path, "image_no_georef.tif")
        label_path = os.path.join(area_path, "label_no_georef.tif")

        if os.path.exists(image_path) and os.path.exists(label_path):
            # 读取原始图像和标签
            image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
            label = cv2.imread(label_path, cv2.IMREAD_UNCHANGED)

            # 去除近红外波段 (去掉第四个波段)
            image = image[:, :, :3]  # 保留前三个波段

            # 确保图像大小可以被512整除，裁剪为512x512
            h, w = image.shape[:2]
            new_h = (h // 512) * 512
            new_w = (w // 512) * 512
            image = image[:new_h, :new_w]
            label = label[:new_h, :new_w]

            # 切割图像为512x512块
            for y in range(0, new_h, 512):
                for x in range(0, new_w, 512):
                    image_patch = image[y:y + 512, x:x + 512]
                    label_patch = label[y:y + 512, x:x + 512]

                    # 保存为PNG格式
                    image_patch_path = os.path.join(save_dir, f"{area}_image_{y}_{x}.png")
                    label_patch_path = os.path.join(save_dir, f"{area}_label_{y}_{x}.png")

                    cv2.imwrite(image_patch_path, image_patch)
                    cv2.imwrite(label_patch_path, label_patch)

                    # 添加到数据列表
                    image_paths.append(image_patch_path)
                    label_paths.append(label_patch_path)

    # 划分数据集
    train_images, temp_images, train_labels, temp_labels = train_test_split(image_paths, label_paths,
                                                                            test_size=(val_ratio + test_ratio),
                                                                            random_state=42)
    val_images, test_images, val_labels, test_labels = train_test_split(temp_images, temp_labels,
                                                                        test_size=test_ratio / (val_ratio + test_ratio),
                                                                        random_state=42)

    # 创建数据集文件夹
    train_dir = os.path.join(save_dir, "train")
    val_dir = os.path.join(save_dir, "val")
    test_dir = os.path.join(save_dir, "test")

    # 创建image和label子文件夹
    os.makedirs(os.path.join(train_dir, "image"), exist_ok=True)
    os.makedirs(os.path.join(train_dir, "label"), exist_ok=True)
    os.makedirs(os.path.join(val_dir, "image"), exist_ok=True)
    os.makedirs(os.path.join(val_dir, "label"), exist_ok=True)
    os.makedirs(os.path.join(test_dir, "image"), exist_ok=True)
    os.makedirs(os.path.join(test_dir, "label"), exist_ok=True)


    # 保存训练集、验证集和测试集数据
    def save_data(image_paths, label_paths, dataset_type):
        print(f"保存{dataset_type}数据集...")
        for image_path, label_path in zip(image_paths, label_paths):
            # 提取文件名并保存
            image_name = os.path.basename(image_path)
            label_name = os.path.basename(label_path)

            # 将图像和标签移动到相应目录
            cv2.imwrite(os.path.join(dataset_type, "image", image_name.replace("_image", "")), cv2.imread(image_path))
            cv2.imwrite(os.path.join(dataset_type, "label", label_name.replace("_label", "")), cv2.imread(label_path))


    # 保存数据集
    save_data(train_images, train_labels, train_dir)
    save_data(val_images, val_labels, val_dir)
    save_data(test_images, test_labels, test_dir)

print("数据集划分和保存完成！")
